from PIL import Image
from kmeans_pytorch import kmeans
from kmeans_pytorch import kmeans_predict

from utils.FileHelper import openRGBImage


import colorsys
def get_dominant_color(image,num_colors=25):
#颜色模式转换，以便输出rgb颜色值
    image = image.convert('RGBA')
#生成缩略图，减少计算量，减小cpu压力
    image.thumbnail((200, 200))
    max_score = 0
    dominant_color = []
    for count, (r, g, b, a) in image.getcolors(image.size[0] * image.size[1]):
        # 跳过纯黑色
        if a == 0:
            continue
        saturation = colorsys.rgb_to_hsv(r / 255.0, g / 255.0, b / 255.0)[1]
        y = min(abs(r * 2104 + g * 4130 + b * 802 + 4096 + 131072) >> 13, 235)
        y = (y - 16.0) / (235 - 16)
        # 忽略高亮色
        if y > 0.9:
            continue
        # Calculate the score, preferring highly saturated colors.
        # Add 0.1 to the saturation so we don't completely ignore grayscale
        # colors by multiplying the count by zero, but still give them a low
        # weight.
        score = (saturation + 0.1) * count
        if score > max_score:
            max_score = score
            dominant_color.append((r, g, b))
    return dominant_color

def get_dominant_colors(image, num_colors=50):
    # small_image = image.resize((200, 200))
    result = image.convert(
        "P", palette=Image.ADAPTIVE, colors=256
    )  # image with only 10 dominating colors

    # Find dominant colors
    palette = result.getpalette()
    color_counts = sorted(result.getcolors(), reverse=True)
    colors = list()

    for i in range(num_colors):
        palette_index = color_counts[i][1]
        dominant_color = palette[palette_index * 3 : palette_index * 3 + 3]
        colors.append(tuple(dominant_color))

    # print(colors)
    return colors

def draw(colors):
    leng = len(colors)
    x = [xi for xi in range(0,leng+1, 1)]
    y = [yi for yi in range(0, leng + 1,1)]
    c = Image.new('RGB',(leng,leng))
    for i in range(leng):
        value = colors[i]
        c.putpixel([x[i],y[i]],value)
    c.show()

def get_color_pallete(img,num_colors,device):
    pass

if __name__ == '__main__':
    import torch
    import torchvision.transforms as T
    img = openRGBImage('./real.png')
    colors = get_dominant_colors(img)
    print(colors)
    draw(colors)
    # t = torch.Tensor(colors) / 255
    # print(t.shape)
    # t = torch.reshape(t,(3,7,7))
    # print(t.shape)
    # trans = T.Compose([T.Normalize((0.5,) * 3, (0.5,) * 3)])
    # print(trans(t))
